package com.et.mserver;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.mina.core.session.AttributeKey;
import org.apache.mina.core.session.DummySession;
import org.apache.mina.core.session.IoEvent;
import org.apache.mina.core.session.IoSession;
import org.apache.mina.filter.executor.OrderedThreadPoolExecutor;

/**
 * A {@link ThreadPoolExecutor} that maintains the order of {@link IoEvent}s.
 * <p/>
 * {@link OrderedThreadPoolExecutor}޸, ͬһSession¼Զռʽִ
 *
 * @author qiu_sheng
 */
public class MailOrderedThreadPoolExecutor extends ThreadPoolExecutor {
	/**
	 * A default value for the initial pool size
	 */
	private static final int DEFAULT_INITIAL_THREAD_POOL_SIZE = 0;

	/**
	 * A default value for the maximum pool size
	 */
	private static final int DEFAULT_MAX_THREAD_POOL = 16;

	/**
	 * A default value for the KeepAlive delay
	 */
	private static final int DEFAULT_KEEP_ALIVE = 30;

	private static final IoSession EXIT_SIGNAL = new DummySession();

	/**
	 * A key stored into the session's attribute for the event tasks being queued
	 */
	private final AttributeKey TASKS_QUEUE = new AttributeKey(getClass(), "tasksQueue");

	private final BlockingQueue<IoSession> waitingSessions = new LinkedBlockingQueue<IoSession>();

	private final Set<Worker> workers = new HashSet<Worker>();

	private volatile int largestPoolSize;
	private final AtomicInteger idleWorkers = new AtomicInteger();

	private long completedTaskCount;
	private volatile boolean shutdown;

	/**
	 * Creates a default ThreadPool, with default values :
	 * - minimum pool size is 0
	 * - maximum pool size is 16
	 * - keepAlive set to 30 seconds
	 * - A default ThreadFactory
	 * - All events are accepted
	 */
	public MailOrderedThreadPoolExecutor() {
		this(DEFAULT_INITIAL_THREAD_POOL_SIZE, DEFAULT_MAX_THREAD_POOL,
				DEFAULT_KEEP_ALIVE, TimeUnit.SECONDS, Executors.defaultThreadFactory());
	}

	/**
	 * Creates a default ThreadPool, with default values :
	 * - minimum pool size is 0
	 * - keepAlive set to 30 seconds
	 * - A default ThreadFactory
	 * - All events are accepted
	 *
	 * @param maximumPoolSize The maximum pool size
	 */
	public MailOrderedThreadPoolExecutor(int maximumPoolSize) {
		this(DEFAULT_INITIAL_THREAD_POOL_SIZE, maximumPoolSize, DEFAULT_KEEP_ALIVE, TimeUnit.SECONDS,
				Executors.defaultThreadFactory());
	}

	/**
	 * Creates a default ThreadPool, with default values :
	 * - keepAlive set to 30 seconds
	 * - A default ThreadFactory
	 * - All events are accepted
	 *
	 * @param corePoolSize    The initial pool sizePoolSize
	 * @param maximumPoolSize The maximum pool size
	 */
	public MailOrderedThreadPoolExecutor(int corePoolSize, int maximumPoolSize) {
		this(corePoolSize, maximumPoolSize, DEFAULT_KEEP_ALIVE, TimeUnit.SECONDS,
				Executors.defaultThreadFactory());
	}

	/**
	 * Creates a default ThreadPool, with default values :
	 * - A default ThreadFactory
	 * - All events are accepted
	 *
	 * @param corePoolSize    The initial pool sizePoolSize
	 * @param maximumPoolSize The maximum pool size
	 * @param keepAliveTime   Default duration for a thread
	 * @param unit            Time unit used for the keepAlive value
	 */
	public MailOrderedThreadPoolExecutor(
			int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit) {
		this(corePoolSize, maximumPoolSize, keepAliveTime, unit,
				Executors.defaultThreadFactory());
	}

	/**
	 * Creates a new instance of a OrderedThreadPoolExecutor.
	 *
	 * @param corePoolSize    The initial pool sizePoolSize
	 * @param maximumPoolSize The maximum pool size
	 * @param keepAliveTime   Default duration for a thread
	 * @param unit            Time unit used for the keepAlive value
	 * @param threadFactory   The factory used to create threads
	 * @param queueHandler    The queue used to store events
	 */
	public MailOrderedThreadPoolExecutor(
			int corePoolSize, int maximumPoolSize,
			long keepAliveTime, TimeUnit unit,
			ThreadFactory threadFactory) {
		// We have to initialize the pool with default values (0 and 1) in order to
		// handle the exception in a better way. We can't add a try {} catch() {}
		// around the super() call.
		super(DEFAULT_INITIAL_THREAD_POOL_SIZE, 1, keepAliveTime, unit,
				new SynchronousQueue<Runnable>(), threadFactory, new AbortPolicy());

		if (corePoolSize < DEFAULT_INITIAL_THREAD_POOL_SIZE) {
			throw new IllegalArgumentException("corePoolSize: " + corePoolSize);
		}

		if ((maximumPoolSize == 0) || (maximumPoolSize < corePoolSize)) {
			throw new IllegalArgumentException("maximumPoolSize: " + maximumPoolSize);
		}

		// Now, we can setup the pool sizes
		super.setCorePoolSize(corePoolSize);
		super.setMaximumPoolSize(maximumPoolSize);
	}


	/**
	 * {@inheritDoc}
	 */
	@Override
	public void setRejectedExecutionHandler(RejectedExecutionHandler handler) {
		// Ignore the request.  It must always be AbortPolicy.
	}

	/**
	 * Add a new thread to execute a task, if needed and possible.
	 * It depends on the current pool size. If it's full, we do nothing.
	 */
	private void addWorker() {
		synchronized (workers) {
			if (workers.size() >= super.getMaximumPoolSize()) {
				return;
			}

			// Create a new worker, and add it to the thread pool
			Worker worker = new Worker();
			Thread thread = getThreadFactory().newThread(worker);

			// As we have added a new thread, it's considered as idle.
			idleWorkers.incrementAndGet();

			// Now, we can start it.
			thread.start();
			workers.add(worker);

			if (workers.size() > largestPoolSize) {
				largestPoolSize = workers.size();
			}
		}
	}

	/**
	 * Add a new Worker only if there are no idle worker.
	 */
	private void addWorkerIfNecessary() {
		if (idleWorkers.get() == 0) {
			synchronized (workers) {
				if (workers.isEmpty() || (idleWorkers.get() == 0)) {
					addWorker();
				}
			}
		}
	}

	private void removeWorker() {
		synchronized (workers) {
			if (workers.size() <= super.getCorePoolSize()) {
				return;
			}
			waitingSessions.offer(EXIT_SIGNAL);
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int getMaximumPoolSize() {
		return super.getMaximumPoolSize();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public void setMaximumPoolSize(int maximumPoolSize) {
		if ((maximumPoolSize <= 0) || (maximumPoolSize < super.getCorePoolSize())) {
			throw new IllegalArgumentException("maximumPoolSize: "
					+ maximumPoolSize);
		}

		synchronized (workers) {
			super.setMaximumPoolSize(maximumPoolSize);
			int difference = workers.size() - maximumPoolSize;
			while (difference > 0) {
				removeWorker();
				--difference;
			}
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public boolean awaitTermination(long timeout, TimeUnit unit)
			throws InterruptedException {

		long deadline = System.currentTimeMillis() + unit.toMillis(timeout);

		synchronized (workers) {
			while (!isTerminated()) {
				long waitTime = deadline - System.currentTimeMillis();
				if (waitTime <= 0) {
					break;
				}

				workers.wait(waitTime);
			}
		}
		return isTerminated();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public boolean isShutdown() {
		return shutdown;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public boolean isTerminated() {
		if (!shutdown) {
			return false;
		}

		synchronized (workers) {
			return workers.isEmpty();
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public void shutdown() {
		if (shutdown) {
			return;
		}

		shutdown = true;

		synchronized (workers) {
			for (int i = workers.size(); i > 0; i--) {
				waitingSessions.offer(EXIT_SIGNAL);
			}
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@SuppressWarnings("unchecked")
	@Override
	public List<Runnable> shutdownNow() {
		shutdown();

		List<Runnable> answer = new ArrayList<Runnable>();
		IoSession session;

		while ((session = waitingSessions.poll()) != null) {
			if (session == EXIT_SIGNAL) {
				waitingSessions.offer(EXIT_SIGNAL);
				Thread.yield(); // Let others take the signal.
				continue;
			}

			Queue<Runnable> tasksQueue = (Queue<Runnable>) session.getAttribute(TASKS_QUEUE);

			synchronized (tasksQueue) {

				for (Runnable task : tasksQueue) {
					answer.add(task);
				}

				tasksQueue.clear();
			}
		}

		return answer;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public void execute(Runnable task) {
		if (shutdown) {
			rejectTask(task);
		}

		// Check that it's a IoEvent task
		checkTaskType(task);

		IoEvent event = (IoEvent) task;
		IoSession session = event.getSession();

		// Get the session's queue of events
		Queue<Runnable> tasksQueue = getTasksQueue(session);
		boolean offerSession;
		boolean offerEvent = true;

		if (offerEvent) {
			// Ok, the message has been accepted
			synchronized (tasksQueue) {
				offerSession = tasksQueue.isEmpty();

				// Inject the event into the executor taskQueue
				tasksQueue.offer(event);
			}
		} else {
			offerSession = false;
		}

		if (offerSession) {
			waitingSessions.offer(session);
		}

		addWorkerIfNecessary();
	}

	private void rejectTask(Runnable task) {
		getRejectedExecutionHandler().rejectedExecution(task, this);
	}

	private void checkTaskType(Runnable task) {
		if (!(task instanceof IoEvent)) {
			throw new IllegalArgumentException("task must be an IoEvent or its subclass.");
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int getActiveCount() {
		synchronized (workers) {
			return workers.size() - idleWorkers.get();
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public long getCompletedTaskCount() {
		synchronized (workers) {
			long answer = completedTaskCount;
			for (Worker w : workers) {
				answer += w.completedTaskCount;
			}

			return answer;
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int getLargestPoolSize() {
		return largestPoolSize;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int getPoolSize() {
		synchronized (workers) {
			return workers.size();
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public long getTaskCount() {
		return getCompletedTaskCount();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public boolean isTerminating() {
		synchronized (workers) {
			return isShutdown() && !isTerminated();
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int prestartAllCoreThreads() {
		int answer = 0;
		synchronized (workers) {
			for (int i = super.getCorePoolSize() - workers.size(); i > 0; i--) {
				addWorker();
				answer++;
			}
		}
		return answer;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public boolean prestartCoreThread() {
		synchronized (workers) {
			if (workers.size() < super.getCorePoolSize()) {
				addWorker();
				return true;
			} else {
				return false;
			}
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public BlockingQueue<Runnable> getQueue() {
		throw new UnsupportedOperationException();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public void purge() {
		// Nothing to purge in this implementation.
	}

	/**
	 * {@inheritDoc}
	 */
	@SuppressWarnings("unchecked")
	@Override
	public boolean remove(Runnable task) {
		checkTaskType(task);
		IoEvent event = (IoEvent) task;
		IoSession session = event.getSession();
		Queue<Runnable> tasksQueue = (Queue<Runnable>) session.getAttribute(TASKS_QUEUE);

		if (tasksQueue == null) {
			return false;
		}

		boolean removed;

		synchronized (tasksQueue) {
			removed = tasksQueue.remove(task);
		}

		return removed;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int getCorePoolSize() {
		return super.getCorePoolSize();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public void setCorePoolSize(int corePoolSize) {
		if (corePoolSize < 0) {
			throw new IllegalArgumentException("corePoolSize: " + corePoolSize);
		}
		if (corePoolSize > super.getMaximumPoolSize()) {
			throw new IllegalArgumentException("corePoolSize exceeds maximumPoolSize");
		}

		synchronized (workers) {
			if (super.getCorePoolSize() > corePoolSize) {
				for (int i = super.getCorePoolSize() - corePoolSize; i > 0; i--) {
					removeWorker();
				}
			}
			super.setCorePoolSize(corePoolSize);
		}
	}

	@SuppressWarnings("unchecked")
	private Queue<Runnable> getTasksQueue(IoSession session) {
		Queue<Runnable> tasksQueue = (Queue<Runnable>) session.getAttribute(TASKS_QUEUE);

		if (tasksQueue == null) {
			tasksQueue = new ConcurrentLinkedQueue<Runnable>();
			Queue<Runnable> oldTasksQueue = (Queue<Runnable>) session.setAttributeIfAbsent(TASKS_QUEUE, tasksQueue);

			if (oldTasksQueue != null) {
				tasksQueue = oldTasksQueue;
			}
		}

		return tasksQueue;
	}

	private class Worker implements Runnable {

		private volatile long completedTaskCount;
		private Thread thread;

		public void run() {
			thread = Thread.currentThread();

			try {
				for (; ; ) {
					IoSession session = fetchSession();

					idleWorkers.decrementAndGet();

					if (session == null) {
						synchronized (workers) {
							if (workers.size() > getCorePoolSize()) {
								// Remove now to prevent duplicate exit.
								workers.remove(this);
								break;
							}
						}
					}

					if (session == EXIT_SIGNAL) {
						break;
					}

					try {
						if (session != null) {
							runTasks(getTasksQueue(session));
						}
					} finally {
						idleWorkers.incrementAndGet();
					}
				}
			} finally {
				synchronized (workers) {
					workers.remove(this);
					MailOrderedThreadPoolExecutor.this.completedTaskCount += completedTaskCount;
					workers.notifyAll();
				}
			}
		}

		private IoSession fetchSession() {
			IoSession session = null;
			long currentTime = System.currentTimeMillis();
			long deadline = currentTime + getKeepAliveTime(TimeUnit.MILLISECONDS);
			for (; ; ) {
				try {
					long waitTime = deadline - currentTime;
					if (waitTime <= 0) {
						break;
					}

					try {
						session = waitingSessions.poll(waitTime, TimeUnit.MILLISECONDS);
						break;
					} finally {
						if (session == null) {
							currentTime = System.currentTimeMillis();
						}
					}
				} catch (InterruptedException e) {
					// Ignore.
					continue;
				}
			}
			return session;
		}

		private void runTasks(Queue<Runnable> tasksQueue) {
			for (; ; ) {
				Runnable task;

				synchronized (tasksQueue) {
					if (tasksQueue.isEmpty()) {
						break;
					}

					task = tasksQueue.peek(); // ֱȥTask, ֹµ¼߳ʹñSession

					if (task == null) {
						break;
					}
				}

				try {
					runTask(task);
				} finally {
					synchronized (tasksQueue) {
						task = tasksQueue.poll(); // ִн, ȥTask
					}
				}
			}
		}

		private void runTask(Runnable task) {
			beforeExecute(thread, task);
			boolean ran = false;
			try {
				task.run();
				ran = true;
				afterExecute(task, null);
				completedTaskCount++;
			} catch (RuntimeException e) {
				if (!ran) {
					afterExecute(task, e);
				}
				throw e;
			}
		}
	}
}
